use mcp_daemon::types::Implementation;
use mcp_daemon::sse::http_server::{run_http_server, ServerConfig, CorsConfig};
use mcp_daemon::server::{Server, McpServer};
use mcp_daemon::transport::ServerHttpTransport;
use serde_json::json;
use anyhow::Result;
use futures::FutureExt;
use sqlx::mysql::MySqlPool;
use sqlx::{Row, Column};
use std::sync::Arc;

// 创建服务器
fn build_server(_transport: ServerHttpTransport) -> Result<Server> {
    let server_info = Implementation {
        name: "Example MCP Server".to_string(),
        version: "0.1.0".to_string(),
    };
    
    // 创建MCP服务器
    let mut mcp_server = McpServer::new(server_info.clone());
    
    // 创建MySQL连接池
    let mysql_pool = tokio::runtime::Handle::current().block_on(async {
        MySqlPool::connect("mysql://tynr426:tynr426@203.25.209.130/ht_b2c_test").await
    }).expect("无法连接到MySQL数据库");
    
    let pool = Arc::new(mysql_pool);
    
    // 注册Hello World工具
    let tool_builder = mcp_server.tool_builder("hello_world")
        .description("A simple hello world tool")
        .input_schema(json!({
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "The name to greet"
                }
            },
            "required": ["name"]
        }));
    
    let (tool_metadata, registered_tool) = tool_builder.build(move |args: Option<serde_json::Value>| {
        async move {
            let name = match &args {
                Some(value) => value.get("name").and_then(|v| v.as_str()).unwrap_or("World"),
                None => "World",
            };
            
            mcp_daemon::types::CallToolResponse {
                content: vec![mcp_daemon::types::Content::Text {
                    text: format!("Hello, {}!", name),
                }],
                is_error: None,
                meta: None,
            }
        }.boxed()
    });
    
    mcp_server.register_tool(tool_metadata, registered_tool);
    
    // 注册MySQL查询工具
    let pool_clone = Arc::clone(&pool);
    let mysql_tool_builder = mcp_server.tool_builder("mysql_query")
        .description("Execute a MySQL query")
        .input_schema(json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "SQL查询语句"
                },
                "params": {
                    "type": "array",
                    "description": "查询参数",
                    "items": {
                        "type": "string"
                    }
                }
            },
            "required": ["query"]
        }));
    
    let (mysql_tool_metadata, mysql_registered_tool) = mysql_tool_builder.build(move |args: Option<serde_json::Value>| {
        let pool = Arc::clone(&pool_clone);
        
        async move {
            match &args {
                Some(value) => {
                    let query = match value.get("query").and_then(|v| v.as_str()) {
                        Some(q) => q,
                        None => return mcp_daemon::types::CallToolResponse {
                            content: vec![mcp_daemon::types::Content::Text {
                                text: "缺少查询语句".to_string(),
                            }],
                            is_error: Some(true),
                            meta: None,
                        }
                    };
                    
                    // 为安全起见，只允许SELECT查询
                    if !query.trim().to_lowercase().starts_with("select") {
                        return mcp_daemon::types::CallToolResponse {
                            content: vec![mcp_daemon::types::Content::Text {
                                text: "只允许SELECT查询".to_string(),
                            }],
                            is_error: Some(true),
                            meta: None,
                        };
                    }
                    
                    // 执行查询
                    match sqlx::query(query)
                        .fetch_all(&*pool)
                        .await {
                            Ok(rows) => {
                                // 将结果转换为JSON
                                let mut results = Vec::new();
                                for row in rows {
                                    let columns = row.columns();
                                    let mut row_data = serde_json::Map::new();
                                    
                                    for (i, column) in columns.iter().enumerate() {
                                        let column_name = column.name();
                                        let value = row.try_get::<serde_json::Value, _>(i).unwrap_or(serde_json::Value::Null);
                                        row_data.insert(column_name.to_string(), value);
                                    }
                                    
                                    results.push(serde_json::Value::Object(row_data));
                                }
                                
                                mcp_daemon::types::CallToolResponse {
                                    content: vec![mcp_daemon::types::Content::Text {
                                        text: format!("查询结果: {}", serde_json::to_string_pretty(&results).unwrap()),
                                    }],
                                    is_error: None,
                                    meta: None,
                                }
                            },
                            Err(e) => mcp_daemon::types::CallToolResponse {
                                content: vec![mcp_daemon::types::Content::Text {
                                    text: format!("查询错误: {}", e),
                                }],
                                is_error: Some(true),
                                meta: None,
                            }
                        }
                },
                None => mcp_daemon::types::CallToolResponse {
                    content: vec![mcp_daemon::types::Content::Text {
                        text: "缺少查询参数".to_string(),
                    }],
                    is_error: Some(true),
                    meta: None,
                }
            }
        }.boxed()
    });
    
    mcp_server.register_tool(mysql_tool_metadata, mysql_registered_tool);
    
    // 返回基础服务器
    Ok(mcp_server.server)
}

#[tokio::main]
async fn main() -> Result<()> {
    // 设置服务器配置
    let config = ServerConfig {
        port: 3004,
        cors: Some(CorsConfig {
            allowed_origin: "http://localhost:3000".to_string(),
            allow_credentials: true,
            max_age: Some(3600),
        }),
        tls: None,
    };
    
    println!("Starting MCP server on http://127.0.0.1:3004");
    println!("WebSocket endpoint: ws://127.0.0.1:3004/ws");
    println!("SSE endpoint: http://127.0.0.1:3004/sse");
    
    // 运行HTTP服务器
    run_http_server(config, None, |transport| async move {
        let server = build_server(transport)?;
        Ok(server)
    })
    .await?;
    
    Ok(())
} 